package kodkod.analysis.visitor;

import kodkod.analysis.Problem;
import kodkod.analysis.ordering.DefaultOrdering;
import kodkod.ast.*;
import kodkod.ast.operator.ExprOperator;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.PrintWriter;
import java.util.*;

public class ClusterLeafsByFormula extends AbstractAnalyzer {
	
	Formula currentFormula;
	public HashMap<Formula, ArrayList<String>> clusters = new HashMap<Formula, ArrayList<String>>();
	
	ArrayList<String> leafs = new ArrayList<String>();
	
	ArrayList<String> reflexiveLeafs = new ArrayList<String>();
	ArrayList<String> predicateLeafs = new ArrayList<String>();
	
	boolean reflexiveClosure = false;
	boolean relationPredicate = false;
	
	Set<Relation> relations; 
	
	public ClusterLeafsByFormula(Problem problem) {
		this.relations = problem.getTranslation(new DefaultOrdering()).getPrimaryVarUsage().keySet();
		this.visit(problem.getFormula());
	}
	
	private String getLeafName(LeafExpression leaf) {
		String leaf_name = leaf.name();
		
		if (leaf instanceof Variable) {
			leaf_name = "$" + leaf_name;
		}
		
		return leaf_name;
	}

	@Override
	void beginChild(Node parent, Node child) {
		if (parent instanceof Formula) {
			currentFormula = (Formula)parent;
		}
		if (child instanceof UnaryExpression) {
			if (((UnaryExpression)child).op().equals(ExprOperator.CLOSURE)) {
				reflexiveClosure = true;
			}
		}
		if (child instanceof RelationPredicate) {
			relationPredicate = true;
		}
	}
	
	@Override
	void endChild(Node parent, Node child) {
		if (child instanceof UnaryExpression) {
			if (((UnaryExpression)child).op().equals(ExprOperator.CLOSURE)) {
				reflexiveClosure = false;
			}			
		}
		if (child instanceof RelationPredicate) {
			relationPredicate = false;
		}
	}

	@Override
	void foundLeafExpression(LeafExpression leaf) {
		String leaf_name = getLeafName(leaf);
		
		ArrayList<String> cluster = clusters.get(currentFormula);
		if (cluster == null) {
			cluster = new ArrayList<String>();
		}
		if (!cluster.contains(leaf_name)) {
			cluster.add(leaf_name);
		}
		clusters.put(currentFormula, cluster);
				
		if (!leafs.contains(leaf_name)) {
			leafs.add(leaf_name);
		}
		
		if (reflexiveClosure && !reflexiveLeafs.contains(leaf_name)) {
			reflexiveLeafs.add(leaf_name);
		}
		
		if (relationPredicate && !predicateLeafs.contains(leaf_name)) {
			predicateLeafs.add(leaf_name);
		}
	}

	@Override
	boolean visited(Node node) {
		return false;
	}
	
	
	public void createVisualization(File file) {
		PrintWriter out = null;
		try {
			out = new PrintWriter(file);
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		}
		out.write("digraph TreeClusters {\n");
		
		
		/*
		 * Enumerate Relations and Formulas
		 */
		HashMap<String, Integer> relationIds = new HashMap<String, Integer>();
		HashMap<Formula, Integer> formulaIds = new HashMap<Formula, Integer>();
		
		int id = 0;				
		for (LeafExpression leaf : relations) {
			relationIds.put(leaf.name(), id++);
		}		
		id = 0;
		for (Formula formula : clusters.keySet()) {
			formulaIds.put(formula, id++);
		}
		
		
		/*
		 * Invert Formula <-> Leaf Relation		 
		 */		
		HashMap<String, ArrayList<Formula>> invertedClusters = new HashMap<String, ArrayList<Formula>>();
		
		for (Formula formula : clusters.keySet()) {
			ArrayList<String> cluster = clusters.get(formula);
			for (String leaf_name : cluster) {
				ArrayList<Formula> list = invertedClusters.get(leaf_name);
				if (list == null) list = new ArrayList<Formula>();
				list.add(formula);
				invertedClusters.put(leaf_name, list);
			}
		}
		
		
		HashSet<UndirectedEdge> edges = new HashSet<UndirectedEdge>();
		HashMap<UndirectedEdge, Integer> edgeWeight = new HashMap<UndirectedEdge, Integer>();
		
		for (String leaf1 : invertedClusters.keySet()) {
			for (Formula formula : invertedClusters.get(leaf1)) {
				for (String leaf2 : clusters.get(formula)) {
					UndirectedEdge edge = new UndirectedEdge(leaf1, leaf2);
					if (edges.add(edge)) {
						edgeWeight.put(edge, 1);
					} else {
						int w = edgeWeight.get(edge);
						edgeWeight.put(edge, w+1);
					}
				}
			}
		}
		
		for (LeafExpression leaf : relations) {
			Integer leaf_id = relationIds.get(leaf.name());
			String name = "\"Leaf " + leaf_id + "\\n" + leaf.name() +"\"";
						
			out.write(name);
			
			out.write("[");
			if (reflexiveLeafs.contains(leaf.name()) || reflexiveLeafs.contains("$" + leaf.name())) {
				out.write("color=red,");
			}
			
			if (predicateLeafs.contains(leaf.name()) || predicateLeafs.contains("$" + leaf.name())) {
				out.write("fillcolor=green,");
			}
			out.write("fontcolor=red]");
			
			out.write("\n");
		}
		
		for (UndirectedEdge edge : edges) {
			Integer leaf_id1 = relationIds.get(edge.a);
			Integer leaf_id2 = relationIds.get(edge.b);
			String name1 = "\"Leaf " + leaf_id1 + "\\n" + edge.a +"\"";
			if (leaf_id1 == null) {
				name1 = "\"" + edge.a + "\"";
			}
			String name2 = "\"Leaf " + leaf_id2 + "\\n" + edge.b +"\"";
			if (leaf_id2 == null) {
				name2 = "\"" + edge.b + "\"";
			}
			
			int weight = edgeWeight.get(edge) / 2;
			
			if (name1.equals(name2) && weight == 1) continue;
			
			String color = "black";
			if (leaf_id1 != null && leaf_id2 != null) color = "red";
			
			if (weight > 1) {
				out.write(name1 + "->" + name2 + "[dir=both, shape=none, color="+color+", label=" + weight + "]\n");
			} else {
				out.write(name1 + "->" + name2 + "[dir=both, shape=none, color="+color+"]\n");
			}
		}
		
		out.write("}");
		out.close();
	}
	
	class UndirectedEdge {
		public String a, b;
		public UndirectedEdge(String a, String b) {
			this.a = a;
			this.b = b;
		}
		
		public boolean equals(Object o) {
			if (o == null) {
				return false;
			}
			if (o instanceof UndirectedEdge) {
				UndirectedEdge edge = (UndirectedEdge)o;
				return a.equals(a) && b.equals(b) || a.equals(b) && b.equals(a);
			} else {
				return false;
			}
		}
		
		public int hashCode() {
			return a.hashCode() + b.hashCode();
		}
	}

}